package com.example.blackredtree;

import lombok.Data;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@Data
public class RedBlackTree<T extends Comparable<T>> {
    private static final boolean RED = false;
    private static final boolean BLACK = true;

    private static final Logger logger = LoggerFactory.getLogger(RedBlackTree.class);

    private class Node {
        T data;
        Node left, right, parent;
        boolean color;

        Node(T data, boolean color, Node parent) {
            this.data = data;
            this.color = color;
            this.parent = parent;
        }
    }

    private Node root, nil;

    public RedBlackTree() {
        nil = new Node(null, BLACK, null);
        root = nil;
    }

    public void insert(T data) {
        logger.info("Inserting: " + data);
        Node newNode = new Node(data, RED, null);
        Node y = nil;
        Node x = root;

        while (x != nil) {
            y = x;
            if (newNode.data.compareTo(x.data) < 0) {
                x = x.left;
            } else {
                x = x.right;
            }
        }

        newNode.parent = y;

        if (y == nil) {
            root = newNode;
        } else if (newNode.data.compareTo(y.data) < 0) {
            y.left = newNode;
        } else {
            y.right = newNode;
        }

        newNode.left = nil;
        newNode.right = nil;
        newNode.color = RED;

        insertFixUp(newNode);
    }

    private void insertFixUp(Node z) {
        logger.info("InsertFixUp for: " + z.data);
        while (z.parent.color == RED) {
            if (z.parent == z.parent.parent.left) {
                Node y = z.parent.parent.right;
                if (y.color == RED) {
                    logger.info("Case 1 - Recoloring");
                    z.parent.color = BLACK;
                    y.color = BLACK;
                    z.parent.parent.color = RED;
                    z = z.parent.parent;
                } else {
                    if (z == z.parent.right) {
                        logger.info("Case 2 - Left Rotate");
                        z = z.parent;
                        leftRotate(z);
                    }
                    logger.info("Case 3 - Right Rotate");
                    z.parent.color = BLACK;
                    z.parent.parent.color = RED;
                    rightRotate(z.parent.parent);
                }
            } else {
                Node y = z.parent.parent.left;
                if (y.color == RED) {
                    logger.info("Case 1 - Recoloring");
                    z.parent.color = BLACK;
                    y.color = BLACK;
                    z.parent.parent.color = RED;
                    z = z.parent.parent;
                } else {
                    if (z == z.parent.left) {
                        logger.info("Case 2 - Right Rotate");
                        z = z.parent;
                        rightRotate(z);
                    }
                    logger.info("Case 3 - Left Rotate");
                    z.parent.color = BLACK;
                    z.parent.parent.color = RED;
                    leftRotate(z.parent.parent);
                }
            }
        }
        root.color = BLACK;
    }

    private void leftRotate(Node x) {
        logger.info("Left Rotate: " + x.data);
        Node y = x.right;
        x.right = y.left;
        if (y.left != nil) {
            y.left.parent = x;
        }
        y.parent = x.parent;
        if (x.parent == nil) {
            root = y;
        } else if (x == x.parent.left) {
            x.parent.left = y;
        } else {
            x.parent.right = y;
        }
        y.left = x;
        x.parent = y;
    }

    private void rightRotate(Node x) {
        logger.info("Right Rotate: " + x.data);
        Node y = x.left;
        x.left = y.right;
        if (y.right != nil) {
            y.right.parent = x;
        }
        y.parent = x.parent;
        if (x.parent == nil) {
            root = y;
        } else if (x == x.parent.right) {
            x.parent.right = y;
        } else {
            x.parent.left = y;
        }
        y.right = x;
        x.parent = y;
    }

    public boolean search(T data) {
        logger.info("Searching: " + data);
        Node current = root;
        while (current != nil) {
            int cmp = data.compareTo(current.data);
            if (cmp < 0) {
                current = current.left;
            } else if (cmp > 0) {
                current = current.right;
            } else {
                return true;
            }
        }
        return false;
    }

    public void inorderTraversal(Node node) {
        if (node != nil) {
            inorderTraversal(node.left);
            System.out.print(node.data + " ");
            inorderTraversal(node.right);
        }
    }

    public void printTree() {
        inorderTraversal(root);
        System.out.println();
    }

}


